# coding: utf-8


import demo_common
from gensim.models.doc2vec import Doc2Vec, LabeledSentence, TaggedDocument

'''
1、官方doc2vec的展示例子，其中有很多评论语 https://github.com/RaRe-Technologies/gensim/blob/develop/docs/notebooks/doc2vec-lee.ipynb
'''

if __name__ == '__main__':
    train_set_raw = demo_common.get_train_set(max_count=1000)
    train_set_tag = [TaggedDocument(line_list, tags=[i]) for i, line_list in enumerate(train_set_raw)]
    doc2vec = Doc2Vec(train_set_tag, min_count=1, window=5, vector_size=100, sample=1e-3,
                      # negative=5,
                      workers=4,
                      epochs=70, dbow_words=1)
    # doc2vec.train(train_set_tag, total_examples=doc2vec.corpus_count, epochs=70)
    # 获取一篇新文档的向量
    inferred_vector = doc2vec.infer_vector(
        ['生孩子', '对于', '大多数', '女性', '来说', '这辈子', '必经', '一件', '事情', '但是', '凡事无绝对', '总有', '那么', '一些', '女人', '怀孕', '这件',
         '事情', '上面', '非常', '困难', '有时候', '绝望', '夫妻', '往往', '这些', '问题', '归结为', '一些', '鬼神', '之说', '也是', '哭笑不得', '先生', '妻子',
         '结婚', '已经', '三四年', '之久', '夫妻', '两人', '结婚', '之前', '想着', '孩子', '因此', '结婚之后', '一直在', '筹备', '备孕', '事情', '可是',
         '一年又一年', '过去了', '妻子', '肚子', '始终', '没有', '反应', '也是', '先生', '妻子', '不行', '辗转', '去了', '不少', '医院', '检查', '没能', '发现',
         '任何', '问题', '医生', '对此', '表示', '放松心情', '这个', '时候', '压力', '大越', '焦虑', '越怀', '不上', '孩子', '夫妻', '两人', '也是', '听从',
         '医生', '建议', '山清水秀', '地方', '旅游', '去了', '没想到', '旅游', '途中', '妻子', '验孕棒', '检验', '发现', '怀上', '可是', '夫妻', '两人',
         '欣喜若狂', '赶忙', '家里', '报喜', '妻子', '这个', '时候', '先生', '昨晚', '已经', '胎梦', '梦见', '一个', '仙女', '她说', '宝宝', '给你', '带来',
         '一定', '要生', '好好爱', '抚养', '长大', '然后', '消失了', '夫妻', '两人', '绝望', '之中', '突获', '曙光', '加上', '两人', '比较', '迷信', '的人',
         '认为', '天上的仙', '孩子', '可是', '好景不长', '怀孕之后', '妻子', '总是', '见红', '着急', '夫妻俩', '去医院', '检查', '医生', '检查', '不到', '胎心',
         '建议', '打掉', '孩子', '两人', '不肯', '相信', '直到', '医生', '他们', '检查', '结果', '夫妻', '两人', '抱头痛哭', '妻子', '这个', '时候', '起来',
         '当初', '做梦', '时候', '没有', '答应', '仙女', '要求', '认为', '肯定', '是因为', '这样', '孩子', '没能', '下来', '万幸', '妻子', '一年之后', '怀上',
         '孩子', '可是', '一年', '之中', '夫妻', '两人', '一直', '活在', '内疚', '痛苦', '当中', '不能自拔', '直到', '怀上', '孩子', '夫妻', '两人', '认为',
         '仙女', '可怜', '夫妻', '人才', '妻子', '怀上', '像是', '这种', '迷信', '想法', '其实', '不少', '对于', '那些', '想要', '孩子', '但是', '一直',
         '不上', '夫妻们', '来说', '最重要的', '去往', '具有', '专业', '资质', '大医院', '听从', '医嘱', '配合', '医生', '正确', '因为', '一些人', '因为',
         '身体', '原因', '的确', '比较', '难以', '怀上', '孩子', '千万', '不要', '希望', '寄托', '那些', '虚无缥缈', '事情', '上面', '本文', '图片', '来源于',
         '网络'])
    for index, sim in doc2vec.docvecs.most_similar([inferred_vector]):
        print(train_set_raw[index], sim)
